import torch
import torch.nn as nn
from torch.autograd import Variable

def parse_stage_config(stage_cfg):

    if isinstance(stage_cfg, int):
        return (stage_cfg,), stage_cfg
    elif isinstance(stage_cfg, tuple) or isinstance(stage_cfg, list):
        return stage_cfg, sum(stage_cfg)
    else:
        raise ValueError("Incorrect STPP config {}".format(stage_cfg))

class StructuredTemporalPyramidPooling(torch.nn.Module):
    """
    This the STPP operator for training. Please see the ICCV paper for more details.
    """
    def __init__(self, feat_dim, standalong_classifier=False, configs=(1, (1,2), 1)):

        super(StructuredTemporalPyramidPooling, self).__init__()
        self.sc = standalong_classifier
        self.feat_dim = feat_dim

        starting_parts, starting_mult = parse_stage_config(configs[0])
        course_parts, course_mult = parse_stage_config(configs[1])
        ending_parts, ending_mult = parse_stage_config(configs[2])

        self.feat_multiplier = starting_mult + course_mult + ending_mult
        self.parts = (starting_parts, course_parts, ending_parts)
        self.norm_num = (starting_mult, course_mult, ending_mult)

    def forward(self, ft, seg_split):

        x1 = seg_split[0]
        x2 = seg_split[1]

        src = ft
        n_sample = src.size()[0]

        def get_stage_stpp(stage_ft, stage_parts, norm_num, scaling):

            stage_stpp = []
            stage_len = stage_ft.size(1)
            for n_part in stage_parts:
                ticks = torch.arange(0, stage_len + 1e-5, stage_len / n_part)
                for i in range(n_part):
                    part_ft = stage_ft[:, int(ticks[i]):int(ticks[i+1]), :].mean(dim=1) / norm_num
                    if scaling is not None:
                        part_ft = part_ft * scaling.resize(n_sample, 1)
                    stage_stpp.append(part_ft)

            return stage_stpp

        feature_parts = []
        feature_parts.extend(get_stage_stpp(src[:, :x1, :], self.parts[0], self.norm_num[0], None))  # starting
        feature_parts.extend(get_stage_stpp(src[:, x1:x2, :], self.parts[1], self.norm_num[1], None))  # course
        feature_parts.extend(get_stage_stpp(src[:, x2:, :], self.parts[2], self.norm_num[2], None))  # ending
        stpp_ft = torch.cat(feature_parts, dim=1)

        if not self.sc:
            return stpp_ft, stpp_ft
        else:
            course_ft = src[:, x1:x2, :].mean(dim=1)
            return course_ft, stpp_ft

    def activity_feat_dim(self):
        if self.sc: return self.feat_dim
        else: return self.feat_dim * self.feat_multiplier

    def completeness_feat_dim(self):
        return self.feat_dim * self.feat_multiplier

class SSN(nn.Module):
    # B x 12 x 2688

    def __init__(self,featuredim,seg_split):

        super(SSN,self).__init__()

        self.featuredim = featuredim
        self.seg_split = seg_split
        self.sptt = StructuredTemporalPyramidPooling(featuredim,configs=(1,(1,2),1))
        self.fc_cls = nn.Linear(self.sptt.activity_feat_dim(),201)
        self.fc_pair = nn.Linear(self.sptt.completeness_feat_dim(),2)

    def forward(self, x):

        corse_ft, stpp_ft = self.sptt(x,self.seg_split)
        pred_cls = self.fc_cls(corse_ft)
        pred_pai = self.fc_pair(stpp_ft)

        return pred_cls,pred_pai

class SSN_OnlyPair(nn.Module):

    def __init__(self,featuredim,seg_split):

        super(SSN_OnlyPair,self).__init__()

        self.featuredim = featuredim
        self.seg_split = seg_split
        self.sptt = StructuredTemporalPyramidPooling(featuredim,configs=(1,(1,2),1))
        self.fc_pair = nn.Linear(self.sptt.completeness_feat_dim(),2)

    def forward(self, x):
        corse_ft, stpp_ft = self.sptt(x,self.seg_split)
        pred_pai = self.fc_pair(stpp_ft)
        return pred_pai

class SSN_OnlyClass(nn.Module):

    def __init__(self,featuredim,seg_split):

        super(SSN_OnlyClass,self).__init__()

        self.featuredim = featuredim
        self.seg_split = seg_split
        self.sptt = StructuredTemporalPyramidPooling(featuredim,configs=(1,(1,2),1))
        self.fc_pair = nn.Linear(self.sptt.completeness_feat_dim(),201)

    def forward(self, x):
        corse_ft, stpp_ft = self.sptt(x,self.seg_split)
        pred_pai = self.fc_pair(stpp_ft)
        return pred_pai

if __name__=='__main__':

    x = torch.randn(8,12,202)
    x = Variable(x)
    model = SSN_OnlyClass(202,[3,9,12])
    out = model(x)
    print(out.shape)

